import numpy as np
import matplotlib.pyplot as plt
from sklearn.datasets import load_iris
from sklearn.metrics import adjusted_rand_score

# 设置中文字体，解决中文显示警告
plt.rcParams["font.family"] = ["SimHei"]
plt.rcParams["axes.unicode_minus"] = False  # 确保负号正常显示

# 欧氏距离和KMeans类定义
def euclidean_distance(x1, x2):
    return np.sqrt(np.sum((x1 - x2) **2))

class KMeans:
    def __init__(self, k=3, max_iter=100, tol=1e-4):
        self.k = k
        self.max_iter = max_iter
        self.tol = tol
        self.centroids = None
        self.labels = None

    def fit(self, X):
        n_samples, n_features = X.shape
        random_idx = np.random.choice(n_samples, self.k, replace=False)
        self.centroids = X[random_idx]

        for _ in range(self.max_iter):
            self.labels = np.array([self._closest_centroid(x) for x in X])
            new_centroids = np.zeros((self.k, n_features))
            for i in range(self.k):
                cluster_points = X[self.labels == i]
                new_centroids[i] = np.mean(cluster_points, axis=0)

            if np.linalg.norm(new_centroids - self.centroids) < self.tol:
                break
            self.centroids = new_centroids

    def _closest_centroid(self, x):
        distances = [euclidean_distance(x, centroid) for centroid in self.centroids]
        return np.argmin(distances)

if __name__ == "__main__":
    iris = load_iris()
    X = iris.data[:, 2:]  # 花瓣长度、花瓣宽度
    y_true = iris.target

    kmeans = KMeans(k=3)
    kmeans.fit(X)

    ari_score = adjusted_rand_score(y_true, kmeans.labels)
    print(f"调整兰德指数（ARI）：{ari_score:.4f}")

    plt.figure(figsize=(8, 6))
    plt.scatter(X[:, 0], X[:, 1], c=kmeans.labels, cmap='viridis', edgecolors='k', label='聚类结果')
    plt.scatter(kmeans.centroids[:, 0], kmeans.centroids[:, 1], s=200, c='red', marker='*', label='质心')
    plt.title(f'k-means聚类Iris数据集结果（k=3, ARI={ari_score:.4f}）')
    plt.xlabel('花瓣长度（cm）')
    plt.ylabel('花瓣宽度（cm）')
    plt.legend()
    plt.show()